# main.py
import os
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats

warnings.filterwarnings('ignore')

# ---------- PERT 工具 ----------
class PertDistribution:
    @staticmethod
    def rvs(a, b, c, size=None, random_state=None):
        alpha = 1 + 4 * (b - a) / (c - a)
        beta = 1 + 4 * (c - b) / (c - a)
        return a + (c - a) * stats.beta.rvs(alpha, beta, size=size, random_state=random_state)

    @staticmethod
    def pdf(x, a, b, c):
        alpha = 1 + 4 * (b - a) / (c - a)
        beta = 1 + 4 * (c - b) / (c - a)
        return stats.beta.pdf((x - a) / (c - a), alpha, beta) / (c - a)

    @staticmethod
    def theoretical_skewness(a, b, c):
        alpha = 1 + 4 * (b - a) / (c - a)
        beta = 1 + 4 * (c - b) / (c - a)
        return float(stats.beta(alpha, beta).stats(moments='s'))

    @staticmethod
    def theoretical_kurtosis(a, b, c):
        alpha = 1 + 4 * (b - a) / (c - a)
        beta = 1 + 4 * (c - b) / (c - a)
        return float(stats.beta(alpha, beta).stats(moments='k'))  # Fisher 峰度

# ---------- 采样主流程 ----------
class SamplingProcessor:
    def __init__(self, config_path):
        self.config = pd.read_csv(config_path)
        self.results = []

    # —— 其余收敛检查、KS 检验、预采样/正式采样函数与原脚本完全一致 ——
    # 为节省篇幅，这里仅列改动过的绘图函数及 run() 调用
    # （如需完整代码请直接复制原脚本对应函数至此）

    # ---------- 生成 4 合 1 验证图（已含 a/b/c/d 标注与统计量） ----------
    def generate_plots(self, samples, floor_name, output_path):
        row = self.config[self.config['floor'] == floor_name].iloc[0]
        a, b, c = row['a'], row['b'], row['c']

        # 计算统计量
        sample_skew = stats.skew(samples)
        sample_kurt = stats.kurtosis(samples, fisher=True)
        theory_skew = PertDistribution.theoretical_skewness(a, b, c)
        theory_kurt = PertDistribution.theoretical_kurtosis(a, b, c)

        # 理论分布对象
        alpha = 1 + 4 * (b - a) / (c - a)
        beta_ = 1 + 4 * (c - b) / (c - a)
        rv = stats.beta(alpha, beta_, loc=a, scale=c - a)

        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        fig.suptitle(f'Floor {floor_name} – Verification of PERT Samples vs. Theoretical Distribution',
                     fontsize=14)

        # (a) Histogram + PDF
        ax = axes[0, 0]
        ax.hist(samples, bins=30, density=True, alpha=0.7, label='Sample Histogram')
        x = np.linspace(a, c, 500)
        ax.plot(x, rv.pdf(x), 'r-', label='Theoretical PDF')
        ax.legend()
        ax.set_title('(a) Distribution Histogram')
        # 右上角文本框
        ax.text(0.98, 0.97,
                f'Skewness: {sample_skew:.3f}\nKurtosis: {sample_kurt:.3f}\n'
                f'(Theory: {theory_skew:.3f}, {theory_kurt:.3f})',
                transform=ax.transAxes, ha='right', va='top',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8))

        # (b) QQ-plot
        ax = axes[0, 1]
        stats.probplot(samples, dist=rv, plot=ax)
        ax.set_title('(b) QQ Plot')

        # (c) CDF 对比
        ax = axes[1, 0]
        sorted_s = np.sort(samples)
        ecdf = np.arange(1, len(sorted_s) + 1) / len(sorted_s)
        ax.step(sorted_s, ecdf, where='post', label='Sample CDF')
        ax.plot(x, rv.cdf(x), 'r-', label='Theoretical CDF')
        ax.legend()
        ax.set_title('(c) CDF Comparison')

        # (d) Box-plot
        ax = axes[1, 1]
        ax.boxplot(samples)
        ax.set_title('(d) Box Plot')

        plt.tight_layout(rect=[0, 0, 1, 0.96])
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        plt.savefig(output_path, dpi=300)
        plt.close()

    # ---------- 其余与原脚本一致，这里省略 ----------
    # 若保留原脚本所有函数，请直接粘贴在 generate_plots 下方即可
    # 为便于一次性复制，以下给出精简版 run() 与主入口

    def run(self):
        from tqdm import tqdm
        pre_results = {}
        for _, row in tqdm(self.config.iterrows(), total=len(self.config), desc='Pre-sampling'):
            n_i, samples, seed = self.pre_sampling(row)
            pre_results[row['floor']] = {'n_i': n_i, 'samples': samples, 'final_seed': seed}

        n_unified = max(v['n_i'] for v in pre_results.values())
        print(f'Unified sample size: {n_unified}')

        for _, row in tqdm(self.config.iterrows(), total=len(self.config), desc='Formal'):
            fl = row['floor']
            if pre_results[fl]['n_i'] == n_unified:
                samp = pre_results[fl]['samples']
                ks_ok, p = self.ks_test(samp, dict(a=row['a'], b=row['b'], c=row['c']))
                conv = True
                seed = pre_results[fl]['final_seed']
            else:
                samp, ks_ok, p, conv, seed = self.formal_sampling(row, n_unified,
                                                                   pre_results[fl]['final_seed'])
            os.makedirs(os.path.dirname(row['samples_output_path']), exist_ok=True)
            pd.DataFrame({'value': samp}).to_csv(row['samples_output_path'], index=False)
            self.generate_plots(samp, fl, row['samples_verification_output_path'])

            m = self.check_convergence(samp, row)
            self.results.append({
                'floor': fl,
                'final_n': n_unified,
                'ks_pvalue': p,
                'mean': m['mean'],
                'skewness': m['skewness'],
                'kurtosis': m['kurtosis'],
                'used_seed': seed,
                'converged': conv
            })

        pd.DataFrame(self.results).to_csv(
            os.path.join(os.path.dirname(self.config.iloc[0]['samples_verification_output_path']),
                         'samples_verification.csv'), index=False)
        print('Done! Check output folders.')

    # ---------- 以下函数与原脚本保持一致（复制即可） ----------
    def check_convergence(self, samples, row):
        return dict(mean=np.mean(samples),
                    var=np.var(samples),
                    median=np.median(samples),
                    skewness=stats.skew(samples),
                    kurtosis=stats.kurtosis(samples, fisher=True))

    def ks_test(self, samples, params):
        theory = PertDistribution.rvs(params['a'], params['b'], params['c'],
                                      size=10000, random_state=42)
        stat, p = stats.ks_2samp(samples, theory)
        return p >= 0.05, p

    def pre_sampling(self, row):
        n, max_n, inc = row['initial_n'], row['max_total_n'], row['increment_n']
        seed = row['random_seed']
        a, b, c = row['a'], row['b'], row['c']
        thr_sk, thr_ku = row['skewness_threshold'], row['kurtosis_threshold']
        while n <= max_n:
            s = PertDistribution.rvs(a, b, c, size=n, random_state=seed)
            m = self.check_convergence(s, row)
            sk_d = abs(m['skewness'] - PertDistribution.theoretical_skewness(a, b, c))
            ku_d = abs(m['kurtosis'] - PertDistribution.theoretical_kurtosis(a, b, c))
            if sk_d <= thr_sk and ku_d <= thr_ku:
                if n > row['initial_n']:
                    prev = PertDistribution.rvs(a, b, c, size=n - inc, random_state=seed)
                    pm = self.check_convergence(prev, row)
                    if abs(m['mean'] - pm['mean']) / abs(pm['mean']) <= 0.005 \
                       and abs(m['var'] - pm['var']) / abs(pm['var']) <= 0.03:
                        return n, s, seed
                else:
                    return n, s, seed
            n += inc
            seed += 1
        return n - inc, s, seed  # fallback

    def formal_sampling(self, row, n_unified, seed):
        a, b, c = row['a'], row['b'], row['c']
        max_try = 100
        thr_sk, thr_ku = row['skewness_threshold'], row['kurtosis_threshold']
        for _ in range(max_try):
            s = PertDistribution.rvs(a, b, c, size=n_unified, random_state=seed)
            m = self.check_convergence(s, row)
            sk_d = abs(m['skewness'] - PertDistribution.theoretical_skewness(a, b, c))
            ku_d = abs(m['kurtosis'] - PertDistribution.theoretical_kurtosis(a, b, c))
            if sk_d > thr_sk or ku_d > thr_ku:
                seed += 1
                continue
            ok, p = self.ks_test(s, dict(a=a, b=b, c=c))
            if ok:
                return s, ok, p, True, seed
            seed += 1
        return s, False, None, False, seed


if __name__ == '__main__':
    SamplingProcessor('sampling_parameters2.csv').run()